use std::{cell::Ref, mem::swap, rc::Rc, vec, fs::File, io::Write};

use crate::{
    global::{clamp, MaterialType, K_INFINITY, M_PI, update_progress},
    object::DrawableObj,
    scene::Scene,
    v2f, v3f,
    vector::{Vector2f, Vector3f},
};

struct HitPayload {
    pub tnear: f32,
    pub index: usize,
    pub uv: Vector2f,
    pub hit_obj: Option<Rc<dyn DrawableObj>>,
}

// impl HitPayload {
//     pub fn new() -> Self {
//         Self {
//             tnear: 0.,
//             index: 0,
//             uv: v2f!(),
//             hit_obj: None,
//         }
//     }
// }

pub struct Renderer {
    scene: Scene,
}

fn deg2rad(deg: f32) -> f32 {
    deg * M_PI / 180.0
}

fn reflect(i: Vector3f, n: Vector3f) -> Vector3f {
    i - 2.0 * i.dot_product(&n) * n
}

fn refract(i: Vector3f, n: Vector3f, ior: f32) -> Vector3f {
    let mut cosi = clamp(-1., 1., i.dot_product(&n));
    let mut etai = 1.;
    let mut etat = ior;
    let mut nn = n;
    if cosi < 0. {
        cosi = -cosi;
    } else {
        swap(&mut etai, &mut etat);
        nn = -nn;
    }

    let eta = etai / etat;
    let k = 1. - eta * eta * (1. - cosi * cosi);
    if k < 0. {
        v3f!(0.)
    } else {
        eta * i + (eta * cosi - k.sqrt()) * nn
    }
}

fn fresnel(i: Vector3f, n: Vector3f, ior: f32) -> f32 {
    let mut cosi = clamp(-1., 1., i.dot_product(&n));
    let mut etai = 1.;
    let mut etat = ior;
    if cosi > 0. {
        swap(&mut etai, &mut etat);
    }
    let sint = etai / etat * f32::max(0., 1. - cosi * cosi).sqrt();

    if sint >= 1. {
        1.
    } else {
        let cost = f32::max(0., 1. - sint * sint).sqrt();
        cosi = cosi.abs();
        let rs = ((etat * cosi) - (etai * cost)) / ((etat * cosi) + (etai * cost));
        let rp = ((etai * cosi) - (etat * cost)) / ((etai * cosi) + (etat * cost));
        (rs * rs + rp * rp) / 2.
    }
}

fn trace(
    orig: Vector3f,
    dir: Vector3f,
    objects: Ref<Vec<Rc<dyn DrawableObj>>>,
) -> Option<HitPayload> {
    let mut tnear = K_INFINITY;
    let mut payload: Option<HitPayload> = None;
    for object in objects.iter() {
        let mut tnear_k = K_INFINITY;
        let mut index_k: u32 = 0;
        let mut uv_k = v2f!(0.);
        if object.intersect(orig, dir, &mut tnear_k, &mut index_k, &mut uv_k) && tnear_k < tnear {
            payload = Some(HitPayload {
                tnear: tnear_k,
                index: index_k as usize,
                uv: uv_k,
                hit_obj: Some(Rc::clone(object)),
            });
            tnear = tnear_k;
        }
    }

    payload
}

fn cast_ray(orig: Vector3f, dir: Vector3f, scene: &Scene, depth: i32) -> Vector3f {
    if depth > scene.max_depth {
        v3f!()
    } else {
        let mut hit_color = scene.background_color;
        if let Some(payload) = trace(orig, dir, scene.get_objects()) {
            let hit_point = orig + dir * payload.tnear;
            let mut n = v3f!();
            let mut st = v2f!();

            if let Some(hit_obj) = payload.hit_obj {
                hit_obj.get_surface_properties(
                    hit_point,
                    dir,
                    payload.index as u32,
                    payload.uv,
                    &mut n,
                    &mut st,
                );

                match hit_obj.get_material_type() {
                    MaterialType::Reflection => {
                        let kr = fresnel(dir, n, hit_obj.get_ior());
                        let reflection_direction = reflect(dir, n).normalize();
                        let reflection_ray_orig = if reflection_direction.dot_product(&n) < 0. {
                            hit_point + n * scene.epsilon
                        } else {
                            hit_point - n * scene.epsilon
                        };
                        hit_color =
                            cast_ray(reflection_ray_orig, reflection_direction, scene, depth + 1)
                                * kr;
                    }

                    MaterialType::ReflectionAndRefraction => {
                        let kr = fresnel(dir, n, hit_obj.get_ior());
                        let reflection_direction = reflect(dir, n).normalize();
                        let reflection_ray_orig = if reflection_direction.dot_product(&n) < 0. {
                            hit_point - n * scene.epsilon
                        } else {
                            hit_point + n * scene.epsilon
                        };
                        let refraction_direction = refract(dir, n, hit_obj.get_ior());
                        let refraction_ray_orig = if refraction_direction.dot_product(&n) < 0. {
                            hit_point - n * scene.epsilon
                        } else {
                            hit_point + n * scene.epsilon
                        };

                        let reflection_color =
                            cast_ray(reflection_ray_orig, reflection_direction, scene, depth + 1);
                        let refraction_color =
                            cast_ray(refraction_ray_orig, refraction_direction, scene, depth + 1);
                        hit_color = reflection_color * kr + refraction_color * (1. - kr);
                    }
                    MaterialType::DiffuseAndGlossy => {
                        let mut light_amt = v3f!();
                        let mut light_specular = v3f!();
                        let shadow_point_orig = if dir.dot_product(&n) < 0. {
                            hit_point + n * scene.epsilon
                        } else {
                            hit_point - n * scene.epsilon
                        };

                        for light in scene.get_lights().iter() {
                            let mut light_dir = light.position - hit_point;

                            let light_distance2 = light_dir.dot_product(&light_dir);
                            light_dir = light_dir.normalize();
                            let l_dot_n = f32::max(0., light_dir.dot_product(&n));


                            let shadow_res =
                                trace(shadow_point_orig, light_dir, scene.get_objects());
                            let in_shadow = if let Some(pld) = shadow_res {
                                pld.tnear * pld.tnear < light_distance2
                            } else {
                                false
                            };

                            // println!("{in_shadow}");

                            light_amt += if in_shadow {
                                v3f!()
                            } else {
                                light.intensity * l_dot_n
                            };

                            let reflection_direction = reflect(-light_dir, n);

                            light_specular += f32::max(0., -reflection_direction.dot_product(&dir))
                                .powf(hit_obj.get_specular_exponent())
                                * light.intensity;
                        }

                        let kd = hit_obj.get_kd();
                        let diffuse_color = hit_obj.eval_diffuse_color(st);
                        let id = light_amt;

                        let ks = hit_obj.get_ks();
                        let is = light_specular;
                        // println!("{}, {}, {}, {}, {}", kd, diffuse_color, id, ks, is);

                        hit_color = kd * diffuse_color * id + ks * is;
                        // println!("{hit_color}");
                    }
                }
            }
        }
        hit_color
    }
}

impl Renderer {
    pub fn new(scene: Scene) -> Self {
        Self { scene }
    }

    pub fn render(&self) {
        // let mut frame_buffer: Vec<Vector3f> = Vec::with_capacity(self.scene.width * self.scene.height);
        let mut frame_buffer:Vec<Vector3f> = vec![v3f!(); self.scene.width * self.scene.height];

        let scale = f32::tan(deg2rad(self.scene.fov * 0.5));
        let image_aspect_ratio = 1.0 * self.scene.width as f32 / self.scene.height as f32;
        println!("scale = {scale}, image_aspect_ratio = {image_aspect_ratio}.");

        let eye_pos = v3f!();
        let mut m: usize = 0;

        for j in 0..self.scene.height {
            for i in 0..self.scene.width {
                let x = 2.0 * scale * image_aspect_ratio / self.scene.width as f32
                    * (i as f32 + 0.5)
                    - scale * image_aspect_ratio;

                let y = -2.0 * scale / self.scene.height as f32 * (j as f32 + 0.5) + scale;

                let dir = v3f!(x, y, -1.).normalize();
                frame_buffer[m] = cast_ray(eye_pos, dir, &self.scene, 0);
                m += 1;
            }
            update_progress(j as f32 / self.scene.height as f32);
        }

        let mut buffer : Vec<u8> = Vec::with_capacity(3 * self.scene.width * self.scene.height);
        for i in frame_buffer {
            buffer.push((255. * clamp(0., 1. , i.x)) as u8);
            buffer.push((255. * clamp(0., 1. , i.y)) as u8);
            buffer.push((255. * clamp(0., 1. , i.z)) as u8);
        }
        
        let mut file = File::create("binary.ppm").expect("Failed to create file");
        let _ = file.write(format!("P6\n{} {}\n255\n", self.scene.width, self.scene.height).as_bytes());
        // for i in 0..self.scene.height * self.scene.width as usize {
        //     let mut buffer = [0u8; 3];
        //     buffer[0] = (255. * clamp(0., 1. , frame_buffer[i].x)) as u8;
        //     buffer[1] = (255. * clamp(0., 1. , frame_buffer[i].y)) as u8;
        //     buffer[2] = (255. * clamp(0., 1. , frame_buffer[i].z)) as u8;
        //     let _ = file.write(&buffer);
        // }
        let _ = file.write(&buffer.as_slice());
        println!("\nFinished.");
    }
}
